import cv
import numpy as np
import scipy.ndimage as ndi

lower = 500; upper = 1000


def identity(array):
  return array[0:2]


def meter_depth(depth):
  meter =  1.0 / (depth * -0.0030711016 + 3.3309495161);
  return meter

def depth_world (coor):
  #assert coor.ndim == 1
  fx_d = 1.0 / 5.9421434211923247e+02
  fy_d = 1.0 / 5.9104053696870778e+02
  cx_d = 3.3930780975300314e+02
  cy_d = 2.4273913761751615e+02

  worldCoor = np.array([(coor[0] - cx_d) * fx_d * coor[2], (coor[1] - cy_d) * fy_d * coor[2], coor[2]])
  return worldCoor 



def world_color (coor):
  #assert coor.ndim == 1

  rotmat = np.array ([[9.9984628826577793e-01, 1.2635359098409581e-03, -1.7487233004436643e-02],
                             [-1.4779096108364480e-03, 9.9992385683542895e-01, -1.2251380107679535e-02],
                              [1.7470421412464927e-02, 1.2275341476520762e-02, 9.9977202419716948e-01]])
  trans = np.array([1.9985242312092553e-02, -7.4423738761617583e-04, -1.0916736334336222e-02])
  #coor = np.dot(rotmat, coor) + trans
  
  
  fx_rgb = 5.2921508098293293e+02
  fy_rgb = 5.2556393630057437e+02
  cx_rgb = 3.2894272028759258e+02
  cy_rgb = 2.6748068171871557e+02
  invZ = 1.0/coor[2]
  result = np.array ([coor[0]*fx_rgb*invZ + cx_rgb, coor[1]*fy_rgb*invZ + cy_rgb])
  return result
  
def matchrgb(depth, rgb):
    # allocate result, r-by-c-by-3
    newrgb = np.empty(rgb.shape)
    
    # get dimensions
    (rows, cols) = depth.shape
    
    # prepare input for world convert, vertex is N-by-3
    vertex = np.empty((rows * cols, 3))
    # x = 0 ... col-1 0 ... col-1 ... repeat row times
    vertex[:,0] = np.tile(np.arange(cols), rows).astype(float)
    # y = 0 0 0 ... repeat col times 1 1 1 ... repeat row times
    vertex[:,1] = np.tile(np.arange(rows), (cols, 1)).T.flatten()
    vertex[:,2] = depth.flatten()
    
    # do world convert, coord is 3-by-N
    coord = depth_world(vertex.T)
    
    # do color convert, rez is 2-by-N
    rez = world_color(coord)
    
    # clip the rez matrix, index is N-by-2
    index = np.empty((rows * cols, 2))
    x = np.round(rez[0,:].clip(0,cols-1))
    y = np.round((rez[1,:]-5).clip(0,rows-1))
    index[:,0] = y.T
    index[:,1] = x.T
        
    index = index.astype(int)
    
    # do color matching
    newrgb = rgb[index[:,0],index[:,1]].reshape(rows, cols, 3)
    return newrgb

# dumb method for conversion
def matchrgb2(depth, rgb):
  newrgb = np.zeros(rgb.shape)
  for r in range(depth.shape[0]):
    for c in range(depth.shape[1]):
      coord = depth_world(np.array([c, r, depth[r,c]]))
      rez = world_color(coord)
      x = int(np.round(rez[0:1].clip(0,639)[0]))
      y = int(np.round((rez[1:2]-5).clip(0,479)[0]))
      newrgb[r,c] = rgb[y,x]
  return newrgb

#def robustsmt(array, sigma = 5):



def robustavg(array):

  if len(array.shape) == 4: # that is, rgb array
    return array.mean(0)

  elif len(array.shape) == 3: # depth array
    weights = ((array > lower ) & (array < upper)).astype(float) + 1e-50
    print weights.shape
    return (array*weights).sum(0) / weights.sum(0)
  else:
    print 'robustavg: invalid array'


def pretty_depth(depth):
  """Converts depth into a 'nicer' format for display

  This is abstracted to allow for experimentation with normalization

  Args:
      depth: A numpy array with 2 bytes per pixel

  Returns:
      A numpy array that has been processed whos datatype is unspecified
  """

  min = 300
  max = 1500
  np.clip(depth, min, max, depth)

  depth -= min
  depth *= 256./(max-min)
  depth = depth.astype(np.uint8)
  return depth

def depth_cv(depth):
  """Converts depth into a 'nicer' format for display

  This is abstracted to allow for experimentation with normalization

  Args:
      depth: A numpy array with 2 bytes per pixel

  Returns:
      An opencv image who's datatype is unspecified
  """
  depth = pretty_depth(depth)
  image = cv.CreateImageHeader((depth.shape[1], depth.shape[0]),
                               cv.IPL_DEPTH_8U,
                               1)
  cv.SetData(image, depth.tostring(),
             depth.dtype.itemsize * depth.shape[1])
  return image


def video_cv(video):
  """Converts video into a BGR format for opencv

  This is abstracted out to allow for experimentation

  Args:
      video: A numpy array with 1 byte per pixel, 3 channels RGB

  Returns:
      An opencv image who's datatype is 1 byte, 3 channel BGR
  """
  video = video[:, :, ::-1]  # RGB -> BGR
  image = cv.CreateImageHeader((video.shape[1], video.shape[0]),
                               cv.IPL_DEPTH_8U,
                               3)
  cv.SetData(image, video.tostring(),
             video.dtype.itemsize * 3 * video.shape[1])
  return image
